#!/usr/bin/env python
import MySQLdb
from pygraphviz import *

class SQL:
    def __init__(self, dname, tname):
        self.conn = MySQLdb.connect(host = "localhost", user = "ssp", \
                passwd = "ssp", db = dname)
        self.cursor = self.conn.cursor()
        self.tn = tname

    def free(self):
        self.cursor.close()
        self.conn.close()

    def tpl_sql(self, cmd = " ", tbc = "*"):
        command = "select " + tbc + " from " + self.tn + " " + cmd
        self.cursor.execute(command)
        result = []
        while(True):
            row = self.cursor.fetchone()
            if row == None:
                break
            result.append(row)
        return result
    
    def ls(self):
        return self.tpl_sql()

    def numbb(self):
        tbc = "MAX(block)"
        return self.tpl_sql(tbc = tbc)[0][0]

    def numnode(self):
        tbc = "MAX(id)"
        return self.tpl_sql(tbc = tbc)[0][0]



class GIMPLE:
    def __init__(self, raw):
        self.id = raw[0]
        self.type = raw[1]
        if(raw[2] == 'NULL'):
            self.op0 = None
        else:
            self.op0 = raw[2]
        if(raw[3] == 'NULL'):
            self.op1 = None
        else:
            self.op1 = raw[3]
        if(raw[4] == 'NULL'):
            self.op2 = None
        else:
            self.op2 = raw[4]
        self.bb = raw[5]
        if(raw[6] == '(null)'):
            self.action = None
        else:
            self.action = raw[6]

    def debug(self):
        print "ID = ",
        print self.id,
        print "\t type = ",
        print self.type,
        print "\t op0 = ",
        print self.op0,
        print "\t op1 = ",
        print self.op1,
        print "\t op2 = ",
        print self.op2,
        print "\t action = ",
        print self.action
        print "\t block = ",
        print self.bb

    def drawbb(self, graph):
        block = self.bb
        subg = graph
        if block:
            if self.type == 'gimple_pred':
                out = "bb"+str(block)+"_entry"
                if self.op0 != None:
                    in1 = "bb" + self.op0 + "_exit"
                else:
                    in1 = None
                if self.op1 != None:
                    in2 = "bb" + self.op1 + "_exit"
                else:
                    in2 = None
                if self.op2 != None:
                    in3 = "bb" + self.op2 + "_exit"
                else:
                    in3 = None

                if in1 != None:
                    subg.add_edge(in1, out)
                if in2 != None:
                    subg.add_edge(in2, out)
                if in3 != None:
                    subg.add_edge(in3, out)
            
            if self.type == 'gimple_succ':
                out = "bb"+str(block)+"_exit"
                if self.op0 != None:
                    in1 = "bb" + self.op0 + "_entry"
                else:
                    in1 = None
                if self.op1 != None:
                    in2 = "bb" + self.op1 + "_entry"
                else:
                    in2 = None
                if self.op2 != None:
                    in3 = "bb" + self.op2 + "_entry"
                else:
                    in3 = None

                if in1 != None:
                    subg.add_edge(out, in1)
                if in2 != None:
                    subg.add_edge(out, in2)
                if in3 != None:
                    subg.add_edge(out, in3)
    def drawhead(self, bb):
        block = self.bb
        subg = bb[self.bb-1][0]
        if block:
            entry = 'bb' + str(block) + '_entry'
            exit = 'bb' + str(block) + '_exit'
            subg.add_node(entry, label = 'entry')
            subg.add_node(exit, label = 'exit')

    def drawtree(self, bb):
        block = self.bb
        prefix = 'bb' + str(block) + '_'
        #prefix = ''
        if self.op0 != None:
            root = prefix + self.op0
        else:
            root = None
        if self.op1 != None and self.op1 != 'CONST':
            left = prefix + self.op1
        else:
            left = None
        if self.op2 != None and self.op2 != 'CONST':
            right = prefix + self.op2
        else:
            right = None
        if self.action != None:
            action = self.action
        else:
            action = ""
        subg = bb[self.bb-1][0]
        if block:
            if True:
                if self.type == 'gimple_assign' or self.type == 'gimple_phi' :
                    if root != None:
                        subg.add_node(root, label = self.op0)
                    if left != None and root:
                        subg.add_node(left, label = self.op1)
                        subg.add_edge(left, root, label = action)
                    if right != None and root:
                        subg.add_node(right, label = self.op2)
                        subg.add_edge(right, root, label = action)
                if False and self.type == 'gimple_array' :
                    if root != None:
                        subg.add_node(root, label = self.op0)
                    if left != None and root:
                        subg.add_node(left, label = self.op1)
                        subg.add_edge(left, root, label = action, style = \
                                'dotted')
                    if right != None and root:
                        subg.add_node(right, label = self.op2)
                        subg.add_edge(right, root, label = action, style = \
                                'dotted')



    def drawit(self, bb):
        label = ""
        shape = 'ellipse'
        if self.type == 'gimple_pred':
            #label = 'bb_start'
            shape = 'point'
        elif self.type == 'gimple_succ':
            #label = 'bb_end'
            shape = 'point'
        elif self.type == 'gimple_assign':
            label = self.op0
            """
            if self.op0.startswith('T.') == False:
		label = self.op0
	    else:
	        label = ''
	        shape = 'point'
            """
        elif self.type == 'gimple_cond':
            label = self.action
        elif self.type == 'gimple_phi':
            label = self.type
        elif self.type == 'gimple_return':
            label = self.type
        else:
            label = self.type
        bb[self.bb-1][0].add_node(self.id, label = label, shape = shape)
        bb[self.bb-1].append(self)

class GRAPH:
    def __init__(self, g):
        self.G = g

    def draw(self):
        self.G.layout(prog = 'dot')
        print 'drawing cfg now ...'
        self.G.draw('test.png')
        self.G.draw('test.dot')

    def make_io(self):
        for sub in self.G.subgraphs_iter():
            for g in sub.iternodes():
                if g.attr['label'] not in ['entry', 'exit', 'Entry', 'Exit']:
                    in_d = self.G.in_degree(g)
                    ou_d = self.G.out_degree(g)
                    if ou_d == 0 and in_d > 0 : # only in, no out
                        print 'OUTPUT', g.attr['label']
                    elif in_d == 0 and ou_d > 0:
                        print 'INPUT', g.attr['label']


    def debug(self):
        self.make_io()
        for g in self.G.iternodes():
            pass

def search(gimple, value):
    for x in gimple:
        for y in x[1:len(x)]:
            if y.op0 == value:
                return y
    return None

G = AGraph(directed = True)
G.graph_attr['label'] = 'ssp_cfg by SSP_Analyzer v0.2 - Xingzhong'
G.graph_attr['rankdir'] = 'LR'
G.graph_attr['clusterrank'] = 'local'
G.graph_attr['compound'] = 'true'
G.graph_attr['remincross'] = 'true'
G.graph_attr['size'] = 20
G.graph_attr['ratio'] = 'auto'
G.graph_attr['ranksep'] = 0.2
#sql = SQL("SSP_GIMPLE_test20", "`work@gr_adaptive_fir_ccf`")
#sql = SQL("SSP_GIMPLE_test19", "`filter@filter`")
sql = SQL("SSP_GIMPLE_test56", "`test@test`")
res = sql.ls()

bb_list = []
#bb_list is the 2-dim list cluster all of the gimple node
#Each element is a basic_block structure
#For each basic_block, the 1st one is the graph object
#rest of them are the gimple object belong to the block
for nb in range(sql.numbb()):
    bb_list.append( [G.subgraph(nbunch=[], name = 'cluster_bb'+ str(nb+1),\
            label = 'bb'+str(nb+1), style = 'filled', color = 'lightblue',\
            )])

NUM = sql.numnode()

X_LIST = []
for x in res:
    X = GIMPLE(x)
    #X.debug()
    #X.drawit(bb_list)
    X.drawhead(bb_list)
    X_LIST.append(X)
G.add_node('bb0_exit', label = 'Entry', style = 'filled', color = 'pink')
G.add_node('bb1_entry', label = 'Exit', style = 'filled', color = 'pink')
for x in X_LIST:
    x.drawbb(G)
for x in X_LIST:
    x.drawtree(bb_list)

"""
G.add_node(0, label='start', style = 'filled', color = 'pink')
G.add_node(NUM+1, label='end', style = 'filled', color = 'pink')

for x in bb_list:
    for y in x[2:len(x)]:
        #G.add_edge(y.id-1, y.id, style = 'bold')
        if y.type == "gimple_succ":
            if y.op0 != None:
                if y.op0 == '1':
                    G.add_edge(y.id, NUM+1)
                else:
                    succ = bb_list[int(y.op0)-1][1]
                    G.add_edge(y.id, succ.id, weight = 10)
                    #print y.id, succ.id
            if y.op1 != None:
                if y.op1 == '1':
                    G.add_edge(y.id, NUM+1)
                else:
                    succ = bb_list[int(y.op1)-1][1]
                    G.add_edge(y.id, succ.id, weight = 10)
                    #print y.id, succ.id
            if y.op2 != None:
                if y.op2 == '1':
                    G.add_edge(y.id, NUM+1)
                else:
                    succ = bb_list[int(y.op2)-1][1]
                    G.add_edge(y.id, succ.id, weight = 10)
                    #print y.id, succ.id
        if y.type == "gimple_assign" or y.type == "gimple_phi":
            label = ""
            label = y.action
            if y.op1 != None:
                gim = search(bb_list, y.op1)
                if gim != None:
                    G.add_edge(y.id, gim.id, style = 'dotted', label = label)
                else:
                    G.add_node(y.op1, style='filled', color = 'yellow')
                    if y.op1 != 'CONST':
                        G.add_edge(y.id, y.op1, style = 'dotted', label = label)
            if y.op2 != None:
                gim = search(bb_list, y.op2)
                if gim != None:
                    G.add_edge(y.id, gim.id, style = 'dotted', label = label)
                else:
                    G.add_node(y.op2, style='filled', color = 'yellow')
                    if y.op2 != 'CONST':
                        G.add_edge(y.id, y.op2 , style = 'dotted', label = label)

        if y.type == "gimple_cond":
            if y.op0 != None:
                gim = search(bb_list, y.op0)
                if gim != None:
                    G.add_edge(y.id, gim.id, style = 'dotted', label = label)
                else:
                    G.add_node(y.op0, style='filled', color = 'yellow')
                    if y.op0 != 'CONST':
                        G.add_edge(y.id, y.op0, style = 'dotted', label = label)
            if y.op1 != None:
                gim = search(bb_list, y.op1)
                if gim != None:
                    G.add_edge(y.id, gim.id, style = 'dotted', label = label)
                else:
                    G.add_node(y.op1, style='filled', color = 'yellow')
                    if y.op1 != 'CONST':
                        G.add_edge(y.id, y.op1, style = 'dotted', label = label)

"""

GG = GRAPH(G)
GG.debug()
GG.draw()
sql.free()
